import string
import random
from menu import new_pause_menu, new_main_menu, new_story_branch_menu
from dialog import Dialog, new_how_to_play_dialog, new_about_dialog, new_game_over_dialog
from data_manager import DataManager
from puzzle import Puzzle, new_puzzle
from word_wrap import wrap_puzzle_text

class GameState:
    def __init__(self, data_manager, user_preferences, puzzle_dimensions):
        self.puzzle_dimensions = puzzle_dimensions
        self.user_preferences = user_preferences
        self.data_manager = data_manager
        self.mode = None
        self.screen = None
        self.dialog_list = []
        self.level_history = None

    def load_level_history(self, level_history):
        self.level_history = level_history

    def load_story_level(self, story, level_id, skip_pre_narrative):
        self.story = story
        level_data = next((level for level in story['levels'] if level['id'] == level_id), None)
        if level_data is None:
            raise ValueError(f"Level with id {level_id} not found")
        self.level = level_data

        if level_data['id'] not in self.level_history['unlocked_levels']:
            self.level_history['unlocked_levels'].append(self.level['id'])


        self.dialog_list.append(Dialog([level_data['name']], self))

        if not skip_pre_narrative and level_data['pre_narrative'] is not None:
            self.dialog_list.append(Dialog(level_data['pre_narrative'], self))

        if level_data['puzzle'] is not None:
            self.puzzle = Puzzle(level_data['puzzle']['quote'], level_data['puzzle']['key'], level_data['puzzle']['letter_mappings'])
            self.screen = 'puzzle'
        else:
            self.on_story_puzzle_completion()

    def open_main_menu(self):
        self.mode = None
        self.menu = new_main_menu(self)
        self.screen = 'menu'

    def open_pause_menu(self):
        self.menu = new_pause_menu(self)
        self.screen = 'menu'

    def open_how_to_play_dialog(self):
        self.dialog_list.append(new_how_to_play_dialog(self))

    def open_about_dialog(self):
        self.dialog_list.append(new_about_dialog(self))

    def open_game_over_dialog(self):
        self.dialog_list.append(new_game_over_dialog(self))
        if self.mode == 'story':
            self.on_story_puzzle_completion()
        else:
            self.start_new_puzzle()

    def on_story_puzzle_completion(self):
        if self.level['id'] not in self.level_history['completed_levels']:
            self.level_history['completed_levels'].append(self.level['id'])

        if self.level['post_narrative'] is not None:
            self.dialog_list.append(Dialog(self.level['post_narrative'], self))
        if self.level['ending'] is not None:
            if self.level['ending'] == 'win':
                self.dialog_list.append(Dialog(['You win!'], self))
            elif self.level['ending'] == 'lose':
                self.dialog_list.append(Dialog(['You lose!'], self))
            else:
                self.dialog_list.append(Dialog(['Stalemate!'], self))
            self.mode = None
            self.story = None
            self.level = None
            self.open_main_menu()
        elif self.level['next_level'] is not None:
            self.load_story_level(self.story, self.level['next_level'], False)
        elif self.level['choices'] is not None:
            self.screen = 'menu'
            self.menu = new_story_branch_menu(self, self.level['choices'])
        else:
            print("This should not be possible.")

    def show_hint(self):
        self.puzzle.show_hint()
        if self.puzzle.is_puzzle_solved():
            self.open_game_over_dialog()

    def start_over(self):
        if self.mode == 'story':
            self.puzzle.start_over(self.level['puzzle']['letter_mappings'])
        else:
            self.puzzle.start_over()

    def load_puzzle(self, puzzle):
        self.puzzle = puzzle

    def start_new_puzzle(self):
        self.puzzle = new_puzzle(self.data_manager, self.user_preferences.quote_db)

    def page_puzzle(self, direction):
        pages_data = self.get_puzzle_output_pages()
        current_page_index = pages_data['cursor_page']
        pages = pages_data['pages']
        max_page = len(pages)
        if len(pages) <= 1:
            return
        if direction == -1:
            if current_page_index == 0:
                next_page_index = max_page - 1
            else:
                next_page_index = current_page_index - 1
        else:
            if current_page_index + 1 == max_page:
                next_page_index = 0
            else:
                next_page_index = current_page_index + 1


        p = pages[next_page_index]
        original_indices = [i['original_index'] for i in p if 'original_index' in i and i['guess_letter'] != ' ']
        if len(original_indices) == 0:
            # In this case, there is no valid cursor index on the next page. The logic to resolve this is complex, so just give up.
            return
        else:
            new_index = min(original_indices)
        self.puzzle.jump_cursor(new_index)

    def get_puzzle_output_pages(self):
        max_width = self.puzzle_dimensions['puzzle_wrap_width']
        max_height = self.puzzle_dimensions['wrap_height']
        guess_text = self.puzzle.apply_guess()
        puzzle_text = self.puzzle.encrypted_text
        red_marks = self.puzzle.get_red_marks()
        x_offset_count = 0
        y_offset_count = 0
        page_offset_count = 0
        pages = []
        current_page = []
        cursor_page = -1
        current_page_index = 0

        for line in wrap_puzzle_text(puzzle_text, max_width):
             if y_offset_count >= max_height:
                     y_offset_count = 0
                     current_page_index += 1
                     pages.append(current_page)
                     current_page = []
             for x_offset_count, puzzle_letter_wrapped in enumerate(line):
                     puzzle_letter = puzzle_letter_wrapped['char']
                     if ('original_index' in puzzle_letter_wrapped):
                         i = puzzle_letter_wrapped['original_index']
                         guess_letter = guess_text[i]
                         is_red_mark = red_marks[i] == 'r'
                         is_cursor = self.puzzle.cursor_position == i
                     else:
                         guess_letter = ' '
                         is_red_mark = False
                         is_cursor = False

                     if is_cursor:
                         cursor_page = current_page_index

                     puzzle_element = {
                         'x_offset': x_offset_count,
                         'y_offset': y_offset_count,
                         'is_cursor': is_cursor,
                         'is_red_mark': is_red_mark,
                         'puzzle_letter': puzzle_letter,
                         'guess_letter': guess_letter
                     }
                     if ('original_index' in puzzle_letter_wrapped):
                         puzzle_element['original_index'] = puzzle_letter_wrapped['original_index']
                     current_page.append(puzzle_element)
             y_offset_count += 1

        if len(current_page) > 0:
            pages.append(current_page)

        return { 'cursor_page': cursor_page, 'pages': pages }

    def update_state(self, button):
        """Updates game state based on button inputs."""
        if self.screen == 'exit_to_os':
            return
        elif len(self.dialog_list) > 0:
            if button == 'a':
                self.dialog_list[0].select()
            elif button == 'b':
                self.dialog_list[0].cancel()
            return
        elif self.screen == 'menu':
            # Menu navigation logic
            if button == 'up':
                self.menu.navigate(-1)
            elif button == 'down':
                self.menu.navigate(1)
            elif button == 'a':
                self.menu.select()
            elif button == 'b':
                self.menu.cancel()
            return

        # Game state logic (only applies when no menu is open)
        guess_changed = False

        if button == 'left':
            self.puzzle.shift_cursor(direction=-1)
        elif button == 'right':
            self.puzzle.shift_cursor(direction=1)
        elif button == 'up':
            self.puzzle.shift_guess(direction=1)
            guess_changed = True
        elif button == 'down':
            self.puzzle.shift_guess(direction=-1)
            guess_changed = True
        elif button == 'a':
            self.page_puzzle(direction=1)
        elif button == 'b':
            self.page_puzzle(direction=-1)
        elif button == 'm':
            self.open_pause_menu()

        if guess_changed and self.puzzle.is_puzzle_solved():
            self.open_game_over_dialog()
